"""
Complete Reproduction Script for Revised Manuscript
Run this script to regenerate all results and figures.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from corrected_model import ExtendedTLSModel
import generate_figures
import calculate_tables

def main():
    print("=" * 70)
    print("REPRODUCING ANALYSIS FOR REVISED MANUSCRIPT")
    print("Version 3.0 - Extended Lindblad Model")
    print("=" * 70)
    
    # Initialize model
    print("\n1. Initializing Extended TLS Model...")
    model = ExtendedTLSModel(Gamma=1.0)
    
    # Define parameters from manuscript
    Omega_values = [0.05, 0.10, 0.20, 0.50]  # MHz
    gamma_phi_values = [0.05, 0.10, 0.20]    # MHz
    gamma_1 = 0.001                          # MHz (fixed)
    
    # Calculate tables
    print("\n2. Calculating Tables 1 and 2...")
    tables = calculate_tables.generate_tables(model, Omega_values, gamma_phi_values, gamma_1)
    
    # Save table data
    tables['table1'].to_csv('table1_data.csv', index=False)
    tables['table2'].to_csv('table2_data.csv', index=False)
    
    print("   Table 1 saved to 'table1_data.csv'")
    print("   Table 2 saved to 'table2_data.csv'")
    
    # Generate figures
    print("\n3. Generating Figures 1 and 2...")
    fig1_data, fig2_data = generate_figures.generate_figures(model, gamma_1)
    
    # Save figure data
    fig1_data.to_csv('figure1_data.csv', index=False)
    fig2_data.to_csv('figure2_data.csv', index=False)
    
    print("   Figure 1 data saved to 'figure1_data.csv'")
    print("   Figure 2 data saved to 'figure2_data.csv'")
    
    # Print summary
    print("\n4. Analysis Summary:")
    print(f"   - {len(Omega_values)} Rabi frequencies: {Omega_values} MHz")
    print(f"   - {len(gamma_phi_values)} dephasing rates: {gamma_phi_values} MHz")
    print(f"   - Fixed relaxation rate: {gamma_1} MHz")
    print(f"   - Coherence times: {1/(0.5*gamma_1 + np.array(gamma_phi_values))} µs")
    
    print("\n5. Verification:")
    # Verify key results
    test_idx = 4  # Ω = 0.10, γ_φ = 0.10
    P_e_test = tables['table1'].iloc[1, 1]  # Should be ~0.4444
    tau_test = tables['table2'].iloc[4, 3]  # Should be ~9.95 µs
    
    print(f"   Sample verification for Ω=0.10, γ_φ=0.10:")
    print(f"   - P_e = {P_e_test:.4f} (expected: ~0.4444)")
    print(f"   - τ = {tau_test:.2f} µs (expected: ~9.95 µs)")
    
    # Check consistency
    if abs(P_e_test - 0.4444) < 0.001 and abs(tau_test - 9.95) < 0.1:
        print("   ✓ Results consistent with manuscript")
    else:
        print("   ⚠ Minor numerical differences detected")
    
    print("\n" + "=" * 70)
    print("ANALYSIS COMPLETE")
    print("All data files have been generated.")
    print("Check the CSV files and run generate_figures.py separately")
    print("to view the plots.")
    print("=" * 70)

if __name__ == "__main__":
    main()